import os
import cv2
import numpy as np
import torch
from flask import Flask, request, jsonify, render_template, url_for, Response
from werkzeug.utils import secure_filename
from ultralytics import YOLO

app = Flask(__name__)
app.config['UPLOAD_FOLDER'] = 'static/uploads'
app.config['ALLOWED_EXTENSIONS'] = {'png', 'jpg', 'jpeg', 'webp'}
app.config['MAX_CONTENT_LENGTH'] = 16 * 1024 * 1024  # 16MB

# 初始化模型
model = YOLO("best.pt")

def allowed_file(filename):
    return '.' in filename and \
           filename.rsplit('.', 1)[1].lower() in app.config['ALLOWED_EXTENSIONS']

def process_image(file_path):
    """处理单张图片并返回结果"""
    results = model(file_path)
    names_dict = results[0].names
    probs = results[0].probs.data.tolist()
    predicted_class = names_dict[np.argmax(probs)]
    confidence = round(max(probs) * 100, 2)
    return predicted_class, confidence

@app.route('/', methods=['GET', 'POST'])
def index():
    if request.method == 'POST':
        # 文件验证逻辑
        if 'file' not in request.files:
            return render_template('index.html', error='请选择文件')
        
        file = request.files['file']
        if file.filename == '':
            return render_template('index.html', error='未选择文件')
        
        if not allowed_file(file.filename):
            return render_template('index.html', error='仅支持jpg/png格式')

        # 保存文件
        filename = secure_filename(file.filename)
        save_path = os.path.join(app.config['UPLOAD_FOLDER'], filename)
        file.save(save_path)

        # 立即生成图片URL用于预览
        # 保存文件后立即生成图片URL
        image_url = url_for('static', filename=f'uploads/{filename}')
        
        # 确保文件保存成功
        if not os.path.exists(save_path):
            return render_template('index.html', error='文件保存失败', image_url=image_url)
        
        try:
            # 获取识别结果
            predicted_class, confidence = process_image(save_path)
            
            # 生成热力图
            heatmap_filename = f'heatmap_{filename}'
            heatmap_path = os.path.join(app.config['UPLOAD_FOLDER'], heatmap_filename)
            
            # 调用heatmap.py生成热力图
            from yolov8_heatmap import yolo_heatmap
            heatmap_generator = yolo_heatmap(
                weight="best.pt",
                device="cuda:0" if torch.cuda.is_available() else "cpu",
                method="GradCAMPlusPlus",
                layer=[4, 6, 8],
                backward_type="class",
                conf_threshold=0.2,
                ratio=0.02,
                show_result=True,
                renormalize=True,
                task="classify",
                img_size=640
            )
            heatmap_generator.process(save_path, heatmap_path)
            
            # 生成热力图URL
            heatmap_url = url_for('static', filename=f'uploads/{heatmap_filename}')
            
            return render_template('index.html',
                                 image_url=image_url,  # 确保返回原始图片URL
                                 heatmap_url=heatmap_url,
                                 predicted_class=predicted_class,
                                 confidence=confidence)
            
        except Exception as e:
            return render_template('index.html', error=f'处理失败: {str(e)}', image_url=image_url)
    
    return render_template('index.html')

@app.route('/batch_upload', methods=['POST'])
def batch_upload():
    """处理批量上传"""
    if 'files' not in request.files:
        return jsonify({"error": "请选择文件"}), 400
    
    files = request.files.getlist('files')
    results = []
    
    for file in files:
        if file and allowed_file(file.filename):
            filename = secure_filename(file.filename)
            save_path = os.path.join(app.config['UPLOAD_FOLDER'], filename)
            file.save(save_path)
            
            predicted_class, confidence = process_image(save_path)
            results.append({
                "filename": filename,
                "image_url": f"/static/uploads/{filename}",
                "predicted_class": predicted_class,
                "confidence": confidence
            })
    
    return jsonify(results)

def generate_frames():
    """生成摄像头视频流"""
    camera = cv2.VideoCapture(0)
    while True:
        success, frame = camera.read()
        if not success:
            break
        else:
            # 实时推理
            results = model(frame)
            annotated_frame = results[0].plot()
            
            ret, buffer = cv2.imencode('.jpg', annotated_frame)
            frame = buffer.tobytes()
            yield (b'--frame\r\n'
                   b'Content-Type: image/jpeg\r\n\r\n' + frame + b'\r\n')

@app.route('/camera')
def camera():
    """摄像头识别页面"""
    return render_template('camera.html')

@app.route('/video_feed')
def video_feed():
    """视频流路由"""
    return Response(generate_frames(),
                    mimetype='multipart/x-mixed-replace; boundary=frame')

@app.errorhandler(413)
def request_entity_too_large(error):
    return render_template('error.html', 
                         error_message="文件过大",
                         error_details="文件大小不能超过16MB"), 413
@app.route('/upload', methods=['POST'])
def upload_file():
    if 'file' not in request.files:
        return jsonify({'success': False, 'message': '未收到文件'}), 400
    
    file = request.files['file']
    if file.filename == '':
        return jsonify({'success': False, 'message': '空文件名'}), 400
    
    try:
        # 保存文件
        filename = secure_filename(file.filename)
        save_path = os.path.join(app.config['UPLOAD_FOLDER'], filename)
        file.save(save_path)
        
        # 处理图片
        predicted_class, confidence = process_image(save_path)
        
        return jsonify({
            'success': True,
            'data': {
                'predictedClass': predicted_class,
                'confidence': confidence,
                'imageUrl': f'/static/uploads/{filename}'
            }
        })
    except Exception as e:
        return jsonify({'success': False, 'message': f'处理失败: {str(e)}'}), 500

if __name__ == '__main__':
    os.makedirs(app.config['UPLOAD_FOLDER'], exist_ok=True)
    app.run(host='0.0.0.0', port=5000, debug=True)


@app.before_request
def log_request():
    print(f"收到请求: {request.method} {request.url}")
